#!/usr/bin/python3/bin/python3
# -*- coding: UTF-8 -*-

import requests
import time
import urllib3
import sys
import os
import re
# import platform
# import ctypes

# def get_free_space_mb(folder):
#     """
#     获取磁盘剩余空间
#     :param folder: 磁盘路径 例如 D:\\
#     :return: 剩余空间 单位 G
#     """
#     if platform.system() == 'Windows':
#         free_bytes = ctypes.c_ulonglong(0)
#         ctypes.windll.kernel32.GetDiskFreeSpaceExW(ctypes.c_wchar_p(folder), None, None, ctypes.pointer(free_bytes))
#         return free_bytes.value / 1024 / 1024 // 1024
#     else:
#         st = os.statvfs(folder)
#         return st.f_bavail * st.f_frsize / 1024 // 1024

def validateTitle(title):
    rstr = r"[\/\\\:\*\?\"\<\>\|\(\)\ ]"  # '/ \ : * ? " < > |'
    new_title = re.sub(rstr, "_", title)  # 替换为下划线
    return new_title


scriptTime = time.strftime("%Y-%m-%d %H-%M-%S", time.localtime())

args = sys.argv[1:]

if len(args) < 1:
    raise ValueError('缺少房间参数')

roomId = args[0]

quality = 4
cookie = None
if len(args) > 1:
    quality = args[1]
if len(args) > 2:
    cookie = args[2]

# 获取房间开播信息
params = {'id': roomId}
headers = {
    'Accept': 'application/json, text/plain, */*',
    'Accept-Encoding': 'gzip, deflate, br',
    'Accept-Language': 'zh-CN,zh;q=0.9',
    'Host': 'api.live.bilibili.com',
    'Origin': 'https://live.bilibili.com',
    'Referer': 'https://live.bilibili.com/' + roomId,
    'Connection': 'keep-alive',
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/80.0.3987.163 Safari/537.36'
}
if cookie is not None:
    headers['Cookie'] = cookie

r = requests.get('https://api.live.bilibili.com/room/v1/Room/get_info', params=params, headers=headers).json()

# 房间标题
title = r['data']['title']

if r['code'] != 0:
    #    raise Exception('接口返回错误:'+r['message'])
    print('%s :接口返回错误：%s\n' % (scriptTime, r['message']))
    exit(0)

if r['data']['live_status'] != 1:
    #    raise Exception('主播未开播')
    print('%s :主播未开播\n' % (scriptTime))
    exit(0)

roomId = r['data']['room_id']

# 纪录进程PID
pid = os.getpid()

if not os.path.exists('pid'):
    os.makedirs('pid')

fpid = open(file='pid/' + roomId, mode='w')
fpid.write(pid.__str__())
fpid.close()

params = {'cid': roomId, 'quality': quality, 'platform': 'web'}
r = requests.get('https://api.live.bilibili.com/room/v1/Room/playUrl', params=params, headers=headers).json()

if r['code'] != 0:
    # raise Exception('获取直播地址失败:'+r['message'])
    print('%s :获取直播地址失败：%s\n' % (scriptTime, r['message']))
    exit(0)

liveUrl = r['data']['durl'][0]['url']

savedFile = validateTitle(title) + '_' + time.strftime("%Y%m%d%H%M%S", time.localtime()) + '.flv'
bak_savedFile = '/mnt/'+savedFile
chunk_size = 1024 * 1024 * 10

urlInfo = urllib3.util.parse_url(liveUrl)

headers = {
    'Accept': '*/*',
    'Accept-Encoding': 'gzip, deflate',
    'Host': urlInfo.host,
    'Origin': 'https://live.bilibili.com',
    'Reference': 'https://live.bilibili.com',
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/86.0.4240.183 Safari/537.36'
}
if cookie is not None:
    headers['Cookie'] = cookie

s = requests.get(liveUrl, headers=headers, stream=True, allow_redirects=False)

print(s.status_code)

if s.status_code == 301 or s.status_code == 302:
    location = s.headers['Location']
    print(location)
    urlInfo = urllib3.util.parse_url(location)
    headers = {
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9',
        'Accept-Encoding': 'gzip, deflate, br',
        'Accept-Language': 'zh-CN,zh;q=0.9',
        'Connection': 'keep-alive',
        'Referer': 'https://live.bilibili.com/',
        'Origin': 'https://live.bilibili.com',
        'Host': urlInfo.host,
        'Upgrade-Insecure-Requests': '1',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/86.0.4240.183 Safari/537.36'
    }
    if cookie is not None:
        headers['Cookie'] = cookie

    s = requests.get(s.headers['Location'], headers=headers, stream=True)
    print(s.status_code)

if s.status_code == 200:
    with open(savedFile, 'wb+') as fd:
        for chunk in s.iter_content(chunk_size=chunk_size):
            fd.write(chunk)


print('%s :finished' % (scriptTime))
